import time
import torch_geometric.utils
from torch_geometric.data import HeteroData
from torch_geometric.loader import LinkNeighborLoader
import torch
import networkx as nx
import matplotlib.pyplot as plt

import torch_geometric.transforms as T
from torch_geometric.nn import HGTConv, Linear
from origin_data import MyData

from load_data import alarms_df

data = MyData(alarms_df).data
# print(data)
data = T.ToUndirected()(data)


# print(data[("host", "belongsto", "bussiness_tree")])

def descrition_data(_data_dict, _data):
    _descrition_data = {}
    all_list = []
    vdata = _data.edge_label_index.t().tolist()
    # print(vdata)
    for i2j in vdata:
        all_list.append((i2j, _data_dict.all_dict[i2j[0]], _data_dict.all_dict[i2j[1]]))
    return all_list


transform = T.RandomLinkSplit(
    num_val=0.1,
    num_test=0.1,
    disjoint_train_ratio=0.0,
    neg_sampling_ratio=0.0,
    add_negative_train_samples=False,
    is_undirected=True,
    # split_labels=True,
    edge_types=[
        ("alarm", "on", "host"),
        ("alarm", "to", "bussiness_tree"),
        ("host", "belongsto", "bussiness_tree")
    ],
    # rev_edge_types=[
    #     ("host", "rev_on", "alarm"),
    #     ("bussiness_tree", "rev_to", "alarm"),
    #     ("bussiness_tree", "rev_belongsto", "host")
    # ]
)

# train_data, val_data, test_data = transform(data)
print(type(data))
train_data, val_data, test_data = transform(data.to_homogeneous())
print("data........................................................................................................",
      data)
print(
    "train_data ........................................................................................................",
    train_data)
print(
    "val_data ........................................................................................................",
    val_data)
print(
    "test_data ........................................................................................................",
    test_data)
# desc_data = descrition_data(data, train_data)
# # print(desc_data)
# # for i in desc_data:
# #     print(i)
#
# # print("train_data.edge_label_index.to(torch.long)",train_data.edge_label_index.t().tolist())
# # # train_loader
# train_loader = LinkNeighborLoader(
#     data=train_data,
#     num_neighbors=[5, 5],
#     neg_sampling_ratio=2.0,
#     edge_label_index=train_data.edge_label_index.to(torch.long),
#     edge_label=train_data.edge_label.to(torch.long),
#     batch_size=32,
#     shuffle=True
# )
#
#
# for batch in train_loader:
#     print(batch)
# #     print("batch.node_id...........",batch.node_id)
# #     print("batch.node_type...........",batch.node_type)
# #     print("batch.n_id...........", batch.n_id)
# #     print("batch.edge_index...........", batch.edge_index)
# #     print("batch.e_id...........", batch.edge_label_index)
# #     print("batch.edge_label...........", batch.edge_label)
# # #     for i in batch.edge_index.t().tolist():
# # #         print("id1:{} id2:{}, {} #### {}".format(int(batch.n_id[i[0]]),int(batch.n_id[i[1]]),batch.all_dict[int(batch.n_id[i[0]])] ,batch.all_dict[int(batch.n_id[i[1]])]))
# #
# # #     # print([j[0] for j in desc_data])
# # #     data = batch.edge_label_index.t().tolist()
# # #     for i in data:
# # #         print(i)
# # #         print(i in [j[0] for j in desc_data])
# #     break
